#!/usr/bin/env python3
"""
PoXiao大模型训练示例脚本
"""

import os
import sys

# 添加项目根目录到Python路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import torch
from torch.utils.data import Dataset, DataLoader
import argparse

# 创建一个简单的示例数据集
class SimpleDataset(Dataset):
    def __init__(self, size=1000, seq_len=128, vocab_size=1000):
        self.size = size
        self.seq_len = seq_len
        self.vocab_size = vocab_size
        
    def __len__(self):
        return self.size
        
    def __getitem__(self, idx):
        # 生成随机数据
        input_ids = torch.randint(0, self.vocab_size, (self.seq_len,))
        labels = torch.randint(0, self.vocab_size, (self.seq_len,))
        loss_mask = torch.ones(self.seq_len)
        return {
            'input_ids': input_ids,
            'labels': labels,
            'loss_mask': loss_mask
        }

def create_sample_data():
    """创建示例数据文件"""
    data_path = "sample_data.txt"
    with open(data_path, "w") as f:
        for i in range(100):
            f.write(f"这是第{i}个示例句子，用于测试PoXiao大模型训练。\n")
    return data_path

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--use_accelerate", action="store_true", help="使用Accelerate")
    parser.add_argument("--use_torchrun", action="store_true", help="使用TorchRun")
    parser.add_argument("--use_tensorboard", action="store_true", help="使用TensorBoard")
    parser.add_argument("--scheduler", type=str, default="cosine", help="学习率调度器类型: cosine 或 linear")
    args = parser.parse_args()
    
    # 创建示例数据
    data_path = create_sample_data()
    
    # 构建训练命令
    cmd = [
        "python", "scripts/train.py",
        "--data_path", data_path,
        "--epochs", "1",
        "--batch_size", "2",
        "--max_seq_len", "64",
        "--hidden_size", "128",
        "--num_hidden_layers", "2",
        "--log_interval", "10",
        "--save_interval", "50",
        "--output_dir", "./output"
    ]
    
    if args.use_accelerate:
        cmd.append("--use_accelerate")
        
    if args.use_torchrun:
        cmd.append("--use_torchrun")
        
    if args.use_tensorboard:
        cmd.append("--use_tensorboard")
        
    if args.scheduler:
        cmd.extend(["--scheduler_plugin", args.scheduler])
        
    # 执行训练命令
    os.system(" ".join(cmd))
    
    # 清理示例数据
    if os.path.exists(data_path):
        os.remove(data_path)

if __name__ == "__main__":
    main()